#include "stdafx.h"
#include "delayloadUtils.h"
#include <delayimp.h>

static int FilterException(int code, PEXCEPTION_POINTERS ex)
{
	if (code != VcppException(ERROR_SEVERITY_ERROR, ERROR_MOD_NOT_FOUND))
	{
		return EXCEPTION_CONTINUE_SEARCH;
	}

	return EXCEPTION_EXECUTE_HANDLER;
}

bool TryToLoadAllImports(const char* dllName)
{
	DWORD oldMode;
	// Suppress the message-box "The program can't start because xxx.dll is missing from your computer"
	// that will appear if the DLL itself is found, but some dependent DLLs are missing.
	SetThreadErrorMode(SEM_FAILCRITICALERRORS | SEM_NOOPENFILEERRORBOX, &oldMode);
	__try
	{
		return SUCCEEDED(__HrLoadAllImportsForDll(dllName));
	}
	__except (FilterException(GetExceptionCode(), GetExceptionInformation()))
	{
		SetThreadErrorMode(oldMode, NULL);
		return false;
	}

	SetThreadErrorMode(oldMode, NULL);
	return true;
}

static HMODULE TryToLoadDll(LPCSTR szDllName, const wchar_t* const* listOfFolders)
{
	wchar_t sz[MAX_PATH];
	for (int i = 0; listOfFolders[i] != nullptr; ++i)
	{
		const wchar_t* path = listOfFolders[i];
		//_snwprintf_s(sz, ARRAYSIZE(sz), L"%ls\\%hs", path, szDllName);
		SetDllDirectory(path);
		HMODULE hMod = LoadLibraryA(szDllName);
		if (hMod != NULL)
		{
			return hMod;
		}
	}
 	return NULL;
}

static const wchar_t* const* ListOfFolders;

FARPROC WINAPI delayHook(unsigned dliNotify, PDelayLoadInfo pdli)
{
	switch (dliNotify) {
	case dliStartProcessing:

		// If you want to return control to the helper, return 0.
		// Otherwise, return a pointer to a FARPROC helper function
		// that will be used instead, thereby bypassing the rest 
		// of the helper.

		break;

	case dliNotePreLoadLibrary:

		// If you want to return control to the helper, return 0.
		// Otherwise, return your own HMODULE to be used by the 
		// helper instead of having it call LoadLibrary itself.

		break;

	case dliNotePreGetProcAddress:

		// If you want to return control to the helper, return 0.
		// If you choose you may supply your own FARPROC function 
		// address and bypass the helper's call to GetProcAddress.

		break;

	case dliFailLoadLib:

		// LoadLibrary failed.
		// If you don't want to handle this failure yourself, return 0.
		// In this case the helper will raise an exception 
		// (ERROR_MOD_NOT_FOUND) and exit.
		// If you want to handle the failure by loading an alternate 
		// DLL (for example), then return the HMODULE for 
		// the alternate DLL. The helper will continue execution with 
		// this alternate DLL and attempt to find the
		// requested entrypoint via GetProcAddress.
		
		if (pdli->dwLastError == ERROR_MOD_NOT_FOUND)
		{
			return (FARPROC)TryToLoadDll(pdli->szDll, ListOfFolders);;
		}
		
		break;

	case dliFailGetProc:

		// GetProcAddress failed.
		// If you don't want to handle this failure yourself, return 0.
		// In this case the helper will raise an exception 
		// (ERROR_PROC_NOT_FOUND) and exit.
		// If you choose you may handle the failure by returning 
		// an alternate FARPROC function address.


		break;

	case dliNoteEndProcessing:

		// This notification is called after all processing is done. 
		// There is no opportunity for modifying the helper's behavior
		// at this point except by longjmp()/throw()/RaiseException. 
		// No return value is processed.

		break;

	default:

		return NULL;
	}

	return NULL;
}

void InitializeDelayLoadAdditionalDllFoldersHook(std::function<const wchar_t*(void)> funcEnumFolders)
{
	std::vector<const wchar_t*> vec;
	for (;;)
	{
		const wchar_t* sz = funcEnumFolders();
		if (sz == nullptr)
		{
			break;
		}

		wchar_t* copySz = _tcsdup(sz);
		vec.push_back(copySz);
	}
	
	if (vec.size() == 0)
	{
		return;
	}

	const wchar_t** listOfFolders = (const wchar_t**)malloc(sizeof(const wchar_t*)*(1+vec.size()));
	for (int i = 0; i < vec.size(); ++i)
	{
		listOfFolders[i] = vec[i];
	}

	listOfFolders[vec.size()] = nullptr;

	ListOfFolders = listOfFolders;
	__pfnDliFailureHook2 = delayHook;
}